import logging
import paramiko
import time
from datetime import datetime
from typing import Optional, Tuple, List

logger = logging.getLogger(__name__)


class ControllerClient:
    MORE_PROMPTS = [
        "---- More ----",  # 华为设备
        "--More--",  # Cisco设备
        "Press any key to continue",  # 其他设备
    ]
    SYS_PROMPTS = [">", "]", "#", "$"]

    def __init__(
        self,
        name: str,
        ip: str,
        username: str,
        passwd: str,
        port: int = 22,
        timeout: int = 10,
    ):
        """
        SSH客户端控制器 - 支持交互式会话和分页处理

        :param ip: 目标主机IP
        :param username: SSH用户名
        :param passwd: SSH密码
        :param port: SSH端口，默认22
        :param timeout: 连接超时时间(秒)，默认10
        """
        self.client = paramiko.SSHClient()
        self.name = name
        self.ip = ip
        self.port = port
        self.timeout = timeout
        self.connected = False
        self.shell = None

        try:
            # 自动添加主机密钥(生产环境应考虑使用known_hosts)
            self.client.set_missing_host_key_policy(paramiko.AutoAddPolicy())

            logger.info(f"正在连接 {username}@{ip}:{port}...")
            self.client.connect(
                hostname=ip,
                port=port,
                username=username,
                password=passwd,
                timeout=timeout,
                banner_timeout=30,
            )
            self.connected = True
            logger.info(f"成功连接到 {ip}")

            # 创建交互式shell
            self._create_shell()

        except paramiko.AuthenticationException:
            logger.error(f"认证失败，请检查用户名/密码")
            raise
        except paramiko.SSHException as e:
            logger.error(f"SSH连接错误: {str(e)}")
            raise
        except Exception as e:
            logger.error(f"连接 {ip} 时发生意外错误: {str(e)}")
            raise

    def _create_shell(self) -> None:
        """创建交互式shell会话"""
        if not self.connected:
            raise RuntimeError("SSH连接未建立")

        logger.debug("创建交互式shell会话")
        self.shell = self.client.invoke_shell(
            term="vt100",  # 模拟终端类型
            width=160,  # 增加宽度以减少换行
            height=500,  # 增加高度以减少分页
        )

        # 等待shell初始化完成
        time.sleep(1)
        self._read_until_prompt()  # 清除初始输出

    def _check_more(self, output: str) -> bool:
        """检查输出是否包含分页提示"""
        return any(prompt in output for prompt in self.MORE_PROMPTS)

    def _read_until_prompt(self, timeout: int = 10, handle_more: bool = True) -> str:
        """
        读取输出直到检测到提示符或处理完所有分页

        :param timeout: 超时时间(秒)
        :param handle_more: 是否处理分页提示
        :return: 完整输出
        """
        if not self.shell:
            raise RuntimeError("Shell会话未初始化")

        output = ""
        start_time = time.time()

        while True:
            if self.shell.recv_ready():
                chunk = self.shell.recv(65535).decode("utf-8", errors="replace")
                output += chunk

                # 处理分页提示
                if handle_more and self._check_more(chunk):
                    logger.debug("检测到分页提示，发送继续命令")
                    self.shell.send(" ")  # 发送空格继续
                    time.sleep(0.1)  # 短暂等待响应
                    continue

                # 检测常见提示符（根据目标设备修改）
                if (
                    f"{self.name}>" in output
                    or f"{self.name}]" in output
                    or "$" in output
                ):
                    break
            else:
                # 短时间等待，避免CPU占用过高
                time.sleep(0.1)

            # 超时处理
            if time.time() - start_time > timeout:
                logger.warning(f"读取超时，返回已获取内容: {output}")
                break

        return output

    def call(
        self, cmd: str, timeout: Optional[int] = None, expect_prompt: bool = True
    ) -> Tuple[int, List[str], List[str]]:
        """
        执行远程命令（交互式shell模式）

        :param cmd: 要执行的命令
        :param timeout: 命令执行超时时间(秒)，None表示使用默认
        :param expect_prompt: 是否等待命令提示符
        :return: (exit_status, stdout_lines, stderr_lines)
        """
        if not self.connected or not self.shell:
            raise RuntimeError("SSH连接或shell会话未建立")

        logger.debug(f"执行命令: {cmd}")
        start_time = datetime.now()

        # 发送命令
        self.shell.send(f"{cmd}\n")

        # 读取命令输出（处理分页）
        output = self._read_until_prompt(timeout or self.timeout)

        # 解析输出（移除命令本身和提示符）
        lines = output.splitlines()
        if lines and lines[0].startswith(cmd):
            lines = lines[1:]  # 移除命令行本身

        # 移除分页提示行
        clean_lines = []
        for line in lines:
            if not any(prompt in line for prompt in self.MORE_PROMPTS):
                clean_lines.append(line)

        # 模拟退出状态（交互式shell没有明确的退出状态，使用0表示成功）
        exit_status = 0 if not self.shell.recv_stderr_ready() else 1

        duration = (datetime.now() - start_time).total_seconds()
        logger.debug(f"命令执行完成 [状态:{exit_status}, 耗时:{duration:.2f}s]")

        # 简单区分标准输出和错误输出（根据实际设备调整）
        stdout_lines = []
        stderr_lines = []

        for line in clean_lines:
            if line.strip():  # 忽略空行
                if "error" in line.lower() or "failed" in line.lower():
                    stderr_lines.append(line.strip())
                else:
                    stdout_lines.append(line.strip())

        return exit_status, stdout_lines, stderr_lines

    def close(self) -> None:
        """关闭SSH连接"""
        if self.shell:
            try:
                self.shell.close()
                logger.debug("Shell会话已关闭")
            except Exception as e:
                logger.error(f"关闭shell会话时出错: {str(e)}")
            self.shell = None

        if self.connected:
            try:
                self.client.close()
                self.connected = False
                logger.info(f"已断开与 {self.ip} 的连接")
            except Exception as e:
                logger.error(f"关闭连接时出错: {str(e)}")

    def __enter__(self):
        """支持with语句"""
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        """支持with语句自动关闭连接"""
        self.close()


# 使用示例
if __name__ == "__main__":
    logging.basicConfig(
        level=logging.DEBUG,
        format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
    )

    try:
        with ControllerClient("Hexin1", "192.168.16.10", "root", "123456") as ssh:
            # 执行可能产生分页的命令
            status, stdout, stderr = ssh.call("display ip routing-table")

            print(f"Exit Status: {status}")
            print("完整输出:")
            for line in stdout:
                print(f"  {line}")

    except Exception as e:
        print(f"操作失败: {str(e)}")
